package edge

import (
	"context"
	"fmt"

	"github.com/DataDog/KubeHound/pkg/kubehound/graph/adapter"
	"github.com/DataDog/KubeHound/pkg/kubehound/graph/types"
	"github.com/DataDog/KubeHound/pkg/kubehound/graph/vertex"
	"github.com/DataDog/KubeHound/pkg/kubehound/models/converter"
	"github.com/DataDog/KubeHound/pkg/kubehound/storage/cache"
	"github.com/DataDog/KubeHound/pkg/kubehound/storage/storedb"
	"github.com/DataDog/KubeHound/pkg/kubehound/store/collections"
	gremlin "github.com/apache/tinkerpop/gremlin-go/v3/driver"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

func init() {
	Register(&PodPatch{}, RegisterGraphMutation)
}

type PodPatch struct {
	BaseEdge
}

type podPatchGroup struct {
	Role primitive.ObjectID `bson:"_id" json:"role"`
}

func (e *PodPatch) Label() string {
	return "POD_PATCH"
}

func (e *PodPatch) Name() string {
	return "PodPatch"
}

func (e *PodPatch) AttckTechniqueID() AttckTechniqueID {
	return AttckTechniqueContainerAdministrationCommand
}

func (e *PodPatch) AttckTacticID() AttckTacticID {
	return AttckTacticExecution
}

func (e *PodPatch) BatchSize() int {
	if e.cfg.LargeClusterOptimizations {
		// Under optimization this becomes a very cheap operation
		return e.cfg.BatchSize
	}

	return e.cfg.BatchSizeClusterImpact
}

func (e *PodPatch) Processor(ctx context.Context, oic *converter.ObjectIDConverter, entry any) (any, error) {
	typed, ok := entry.(*podPatchGroup)
	if !ok {
		return nil, fmt.Errorf("invalid type passed to processor: %T", entry)
	}

	rid, err := oic.GraphID(ctx, typed.Role.Hex())
	if err != nil {
		return nil, fmt.Errorf("%s edge role id convert: %w", e.Label(), err)
	}

	if e.cfg.LargeClusterOptimizations {
		return map[any]any{
			gremlin.T.Label: vertex.PermissionSetLabel,
			gremlin.T.Id:    rid,
		}, nil
	}

	return rid, nil
}

func (e *PodPatch) Traversal() types.EdgeTraversal {
	return func(source *gremlin.GraphTraversalSource, inserts []any) *gremlin.GraphTraversal {
		g := source.GetGraphTraversal()
		if e.cfg.LargeClusterOptimizations {
			// In large clusters this can explode the number of edges and we can safely assume this is a critical issue
			g.
				Inject(inserts).
				Unfold().
				As("rpp").
				MergeV(__.Select("rpp")).
				Option(gremlin.Merge.OnCreate, __.Fail("missing role vertex on POD_PATCH insert")).
				Option(gremlin.Merge.OnMatch, map[any]any{
					"critical": true,
				}).
				AddE(e.Label()).
				Property("attckTechniqueID", string(e.AttckTechniqueID())).
				Property("attckTacticID", string(e.AttckTacticID())).
				Barrier().Limit(0)
		} else {
			// In smaller clusters we can still show the (large set of) attack paths generated by this attack
			g.V().
				Has("runID", e.runtime.RunID.String()).
				Has("cluster", e.runtime.Cluster.Name).
				Has("class", "Pod").
				As("p").
				V(inserts...).
				Has("critical", false).
				AddE(e.Label()).
				To("p").
				Property("attckTechniqueID", string(e.AttckTechniqueID())).
				Property("attckTacticID", string(e.AttckTacticID())).
				Barrier().Limit(0)
		}

		return g
	}
}

// Stream finds all roles that have pod/patch or equivalent wildcard permissions.
func (e *PodPatch) Stream(ctx context.Context, store storedb.Provider, _ cache.CacheReader,
	callback types.ProcessEntryCallback, complete types.CompleteQueryCallback) error {

	permissionSets := adapter.MongoDB(ctx, store).Collection(collections.PermissionSetName)
	pipeline := []bson.M{
		{
			"$match": bson.M{
				"is_namespaced":        false,
				"runtime.runID":        e.runtime.RunID.String(),
				"runtime.cluster.name": e.runtime.Cluster.Name,
				"rules": bson.M{
					"$elemMatch": bson.M{
						"$and": bson.A{
							bson.M{"$or": bson.A{
								bson.M{"apigroups": ""},
								bson.M{"apigroups": "*"},
							}},
							bson.M{"$or": bson.A{
								bson.M{"resources": "pods"},
								bson.M{"resources": "cronjobs"},
								bson.M{"resources": "daemonsets"},
								bson.M{"resources": "deployments"},
								bson.M{"resources": "jobs"},
								bson.M{"resources": "replicasets"},
								bson.M{"resources": "replicationcontrollers"},
								bson.M{"resources": "statefulsets"},
								bson.M{"resources": "*"},
							}},
							bson.M{"$or": bson.A{
								bson.M{"verbs": "patch"},
								bson.M{"verbs": "update"},
								bson.M{"verbs": "*"},
							}},
							bson.M{"resourcenames": nil}, // TODO: handle resource scope
						},
					},
				},
			},
		},
		{
			"$project": bson.M{
				"_id": 1,
			},
		},
	}

	cur, err := permissionSets.Aggregate(ctx, pipeline)
	if err != nil {
		return err
	}
	defer cur.Close(ctx)

	return adapter.MongoCursorHandler[podPatchGroup](ctx, cur, callback, complete)
}
